#!/usr/bin/env python3
"""
Complete pipeline to process JSONL file with DepthCrafter inference and evaluation
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path


def run_command(cmd, description):
    """Run a command and handle errors"""
    print(f"\n{'='*60}")
    print(f"STEP: {description}")
    print(f"{'='*60}")
    print(f"Running: {' '.join(cmd)}")
    
    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        print("✓ Success!")
        if result.stdout:
            print("STDOUT:", result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        print(f"✗ Error: {e}")
        if e.stdout:
            print("STDOUT:", e.stdout)
        if e.stderr:
            print("STDERR:", e.stderr)
        return False


def main():
    parser = argparse.ArgumentParser(description="Complete DepthCrafter pipeline")
    parser.add_argument("--jsonl_path", type=str, required=True, help="Path to text.jsonl file")
    parser.add_argument("--output_dir", type=str, default="./pipeline_output", help="Output directory")
    parser.add_argument("--skip_inference", action="store_true", help="Skip inference step")
    parser.add_argument("--skip_evaluation", action="store_true", help="Skip evaluation step")
    
    # DepthCrafter parameters
    parser.add_argument("--num_denoising_steps", type=int, default=5, help="Number of denoising steps")
    parser.add_argument("--guidance_scale", type=float, default=1.0, help="Guidance scale")
    parser.add_argument("--max_res", type=int, default=1024, help="Maximum resolution")
    parser.add_argument("--process_length", type=int, default=195, help="Process length")
    parser.add_argument("--cpu_offload", type=str, default="model", help="CPU offload strategy")
    
    args = parser.parse_args()
    
    # Check if JSONL file exists
    if not os.path.exists(args.jsonl_path):
        print(f"Error: JSONL file not found: {args.jsonl_path}")
        sys.exit(1)
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Define paths
    inference_output_dir = os.path.join(args.output_dir, "inference")
    evaluation_output_dir = os.path.join(args.output_dir, "evaluation")
    results_json_path = os.path.join(inference_output_dir, "processing_results.json")
    
    success = True
    
    # Step 1: Run inference
    if not args.skip_inference:
        inference_cmd = [
            "python", "process_jsonl.py",
            "--jsonl_path", args.jsonl_path,
            "--output_dir", inference_output_dir,
            "--num_denoising_steps", str(args.num_denoising_steps),
            "--guidance_scale", str(args.guidance_scale),
            "--max_res", str(args.max_res),
            "--process_length", str(args.process_length),
            "--cpu_offload", args.cpu_offload,
        ]
        
        success = run_command(inference_cmd, "DepthCrafter Inference")
        
        if not success:
            print("Inference failed. Stopping pipeline.")
            sys.exit(1)
    else:
        print("Skipping inference step...")
        if not os.path.exists(results_json_path):
            print(f"Error: Results file not found: {results_json_path}")
            print("Cannot skip inference without existing results.")
            sys.exit(1)
    
    # Step 2: Run evaluation
    if not args.skip_evaluation:
        evaluation_cmd = [
            "python", "evaluate_results.py",
            "--jsonl_path", args.jsonl_path,
            "--results_path", results_json_path,
            "--output_dir", evaluation_output_dir,
        ]
        
        success = run_command(evaluation_cmd, "Evaluation")
        
        if not success:
            print("Evaluation failed.")
            sys.exit(1)
    else:
        print("Skipping evaluation step...")
    
    # Summary
    print(f"\n{'='*60}")
    print("PIPELINE COMPLETE!")
    print(f"{'='*60}")
    print(f"Output directory: {args.output_dir}")
    if not args.skip_inference:
        print(f"Inference results: {inference_output_dir}")
    if not args.skip_evaluation:
        print(f"Evaluation results: {evaluation_output_dir}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()
